# -*- coding: utf-8 -*-
"""
Created on Sun Sep 24 11:09:44 2023

@author: 李小斌
"""
import os
import numpy as np
import csv
import jieba
import jieba.analyse
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence
from gensim.models import KeyedVectors

#设置最终嵌入的维度
dimension=450

def getModel(filename):
    sentences=LineSentence(filename)
    #model = Word2Vec(sentences, vector_size=10, window=5, min_count=1, workers=4)
    model = Word2Vec(sentences, vector_size=dimension, window=5, min_count=1, workers=4)
    model.save("word2vec.model")
    model = Word2Vec.load("word2vec.model")
    model.train(sentences,total_examples=model.corpus_count, epochs=1)
    #word_vectors=model.wv
    #word_vectors.save("word2vec.wordvectors")
    return model

def getVector(model,text):
   #wv = KeyedVectors.load("word2vec.wordvectors", mmap='r')
   keywords=jieba.analyse.tfidf(text)
   total=len(keywords)   
   v_total=np.zeros(dimension)
   if total==0:
       return v_total.tolist()
   for keyword in keywords:
       if model.wv.get_index(keyword,default=-1)!=-1:
       #if wv.get_index(keyword,default=-1)!=-1:
           v=model.wv[keyword]
           v_total=v_total+v
   v_total=v_total/total
   return v_total.tolist()
   
    
def main():
    #model=getModel("embedding.txt")
    print(dimension)
    model=getModel("sgns.weibo.word")
    csv_reader=csv.reader(open("invest01.csv"))
    file01=open('embedding'+str(dimension)+'_01list.csv','w',newline='')     #第1次调研嵌入
    file02=open('embedding'+str(dimension)+'_02list.csv','w',newline='')     #第2次调研嵌入
    file03=open('embedding'+str(dimension)+'_03list.csv','w',newline='')     #第3次调研嵌入
    f_csv01=csv.writer(file01)
    f_csv02=csv.writer(file02)
    f_csv03=csv.writer(file03)
    header=[x for x in range(1,dimension+1)]
    f_csv01.writerow(header)
    f_csv02.writerow(header)
    f_csv03.writerow(header)
    n=0
    for row in csv_reader:
        if n==0:
            n=n+1
            continue
        n=n+1
        text=row[1]
        v=getVector(model,text)
        f_csv01.writerow(v) #第1次调研
        text=row[2]
        v=getVector(model,text)
        f_csv02.writerow(v) #第2次调研
        text=row[3]
        v=getVector(model,text)
        f_csv03.writerow(v) #第3次调研
    file01.close()
    file02.close()
    file03.close()
    
if __name__=="__main__":    
    dimension=450
    main()    
    #os.system('shutdown /s /t 120')  
